#include "kpextDescriptorMgr.h"

using namespace kpext;

DescriptorPool::DescriptorPool(std::shared_ptr<vk::Device> device,
                               const std::vector<uint32_t> &  poolCap,
                               uint32_t maxSetsCount,
                               vk::DescriptorPoolCreateFlags flags
                               ):
    m_descPool(nullptr),
    m_descPoolCFlags(flags),
    m_descSetLayoutRes({})
{
    if(device == nullptr){
        throw std::runtime_error("construct kpext::DescriptorPool with null pointer!");
    }

    m_device = device;
    uint32_t tempui[] = {1, 0, 0, 0, 0};
    uint32_t & storageBufSize = tempui[0];
    uint32_t & uniformTexelBufSize = tempui[1];
    uint32_t & uniformBufSize = tempui[2];
    uint32_t & storageTexelBufSize = tempui[3];
    uint32_t & storageImgSize = tempui[4];
    auto count = poolCap.size();
    for (size_t i = 0; i < count; ++i) {
        if(poolCap[i] > 0) tempui[i] = poolCap[i];
    }

    createDescriptorPool(maxSetsCount,
                         storageBufSize,
                         uniformTexelBufSize,
                         uniformBufSize,
                         storageTexelBufSize,
                         storageImgSize);
}

DescriptorPool::~DescriptorPool(){
    if(!m_device){
        return;
    }

    if(m_descSetLayoutRes.size() > 0){
        for(auto & layout : m_descSetLayoutRes){
            m_device->destroyDescriptorSetLayout(*layout);
            *layout = nullptr;
        }
        m_descSetLayoutRes.clear();
    }

    if(m_descPool){
        resetDescriptorPool();
        m_device->destroyDescriptorPool(*m_descPool);
        m_descPool = nullptr;
    }

    m_device = nullptr;
}

void DescriptorPool::resetDescriptorPool()
{
    if(!m_device || !m_descPool){
        return;
    }
    m_device->resetDescriptorPool(*m_descPool);
}

std::shared_ptr<vk::DescriptorSetLayout>
DescriptorPool::createDescriptorSetLayout(
        const std::vector<vk::DescriptorSetLayoutBinding> &descriptorSetBindings)
{
    if(!m_device) return nullptr;

    if(descriptorSetBindings.size() == 0){
        SPDLOG_ERROR(__FUNCTION__ " input zero size vector!");
    }

    vk::DescriptorSetLayoutCreateInfo descriptorSetLayoutInfo(
                vk::DescriptorSetLayoutCreateFlags(),
                static_cast<uint32_t>(descriptorSetBindings.size()),
                descriptorSetBindings.data());
    auto descSetLayout = std::make_shared<vk::DescriptorSetLayout>();
    auto result = m_device->createDescriptorSetLayout(
                &descriptorSetLayoutInfo, nullptr, descSetLayout.get());
    if(result != vk::Result::eSuccess){
        SPDLOG_ERROR("Create descriptor set layout failed!");
    }
    m_descSetLayoutRes.push_back(descSetLayout);
    return descSetLayout;
}

std::shared_ptr<vk::DescriptorSetLayout> DescriptorPool::createDescriptorSetLayout(
        std::vector<IGetDecriptorInfo *> infoGetters,
        bool reorderBinding,
        uint32_t reorderOffset,
        vk::ShaderStageFlags stage
        )
{
    if(!m_device) return nullptr;

    auto count = infoGetters.size();
    if(count < 1) return nullptr;
    for (uint32_t i = 0; i < count; ++i) {
        if(!infoGetters[i]) return nullptr;
    }

    bool forceStage = true;
    if(stage == vk::ShaderStageFlags(0)){
        forceStage = false;
    }

    std::vector<vk::DescriptorSetLayoutBinding> bindinginfos;
    for (uint32_t i = 0; i < count; ++i) {
        if(reorderBinding) infoGetters[i]->setDescBinding(i + reorderOffset);
        if(forceStage) infoGetters[i]->setDescStage(stage);
        bindinginfos.push_back(
                    infoGetters[i]->getDescriptorSetLayoutBinding()
                    );
    }

    return createDescriptorSetLayout(bindinginfos);
}

void DescriptorPool::destoryDescriptorSetLayout(std::shared_ptr<vk::DescriptorSetLayout> & layout)
{
    if(!m_device) return;
    auto item = std::find(m_descSetLayoutRes.begin(),
                          m_descSetLayoutRes.end(),
                          layout
                          );
    if(item != m_descSetLayoutRes.end()){
        m_device->destroyDescriptorSetLayout(*(*item));
        *item = nullptr;
        m_descSetLayoutRes.erase(item);
        layout = nullptr;
    }
}


DescriptorPool::DescSetPair
DescriptorPool::recreateBufDescSet(uint32_t count,
                                   DescSetPair * descSetPair,
                                   vk::DescriptorType type,
                                   vk::ShaderStageFlags stage
                                   )
{
    if(!m_device || !m_descPool) return {nullptr, nullptr};
    std::shared_ptr<vk::DescriptorSetLayout> result_setlayout(nullptr);
    std::shared_ptr<vk::DescriptorSet> result_set(nullptr);
    if(descSetPair){
        result_set = descSetPair->first;
        result_setlayout = descSetPair->second;
        destoryDescriptorSetLayout(result_setlayout);
    }

    if(count < 1) count = 1;

    { // 重新创建描述符集布局
        if(result_setlayout){
            m_device->destroyDescriptorSetLayout(*result_setlayout);
            result_setlayout = nullptr;
        }

        vk::DescriptorSetLayoutBinding tempbinding(
                    0,
                    type,
                    1, //  binding的变量是数组时使用
                    stage
                    );
        std::vector<vk::DescriptorSetLayoutBinding> bindinginfos;
        for (uint32_t i = 0; i < count; ++i) {
            tempbinding.binding = i;
            bindinginfos.push_back(
                        tempbinding
                        );
        }

        result_setlayout = createDescriptorSetLayout(bindinginfos);
        if(!result_setlayout){
            throw std::runtime_error(__FUNCTION__
                                     " create vk::DescriptorSetLayout failed!");
        }
    }


    if((m_descPoolCFlags & vk::DescriptorPoolCreateFlagBits::eFreeDescriptorSet)
            && result_set
            ){
        m_device->freeDescriptorSets(*m_descPool,
                                     1, result_set.get());
    }

    result_set = allocateDescriptorSet(m_descPool,
                                       result_setlayout);
    if(!result_set){
        throw std::runtime_error(__FUNCTION__
                                 " create vk::DescriptorSet failed!");
    }

    return {result_set, result_setlayout};
}

std::shared_ptr<vk::DescriptorSet> DescriptorPool::recreateDescSet(
        vk::DescriptorSetLayout &layout,
        std::shared_ptr<vk::DescriptorSet> descset)
{
    if(!m_device || !m_descPool) return nullptr;

    if((m_descPoolCFlags & vk::DescriptorPoolCreateFlagBits::eFreeDescriptorSet)
            && descset){
        m_device->freeDescriptorSets(*m_descPool,
                                     1, descset.get());
    }

    descset = std::make_shared<vk::DescriptorSet>();

    vk::DescriptorSetAllocateInfo descriptorSetAllocateInfo(
                *m_descPool,
                1, // 创建多少个set
                & layout);

    auto result = m_device->allocateDescriptorSets(&descriptorSetAllocateInfo,
                                                   descset.get());
    if(result != vk::Result::eSuccess){
        SPDLOG_ERROR(__FUNCTION__
                     " create vk::Descriptor failed!");
        return nullptr;
    }

    return descset;
}

bool DescriptorPool::recreateDescSet(DescSetPair &descpair)
{
    if(!m_device || !m_descPool || !descpair.second) return false;
    auto descset = recreateDescSet(*descpair.second, descpair.first);
    descpair.first = descset;
    return (descset != nullptr);
}

bool DescriptorPool::writeDescriptorSet(const vk::WriteDescriptorSet &writeinfo)
{
    if(!m_device) return false;
    if(!writeinfo.dstSet) return false;
    m_device->updateDescriptorSets(1, &writeinfo,
                                   0, nullptr  //  可从其它集复制写入描述信息
                                   );
    return true;
}

bool DescriptorPool::writeDescriptorSet(const std::vector<vk::WriteDescriptorSet> &writeinfos)
{
    if(!m_device) return false;
    uint32_t count = static_cast<uint32_t>(writeinfos.size());
    if(count < 1) return false;
    for (uint32_t i = 0; i < count; ++i) {
        if(!writeinfos[i].dstSet) return false;
    }
    m_device->updateDescriptorSets(writeinfos, {});
    return true;
}

bool DescriptorPool::writeDescriptorSet(
        std::shared_ptr<vk::DescriptorSet> dstSet,
        std::vector<IGetDecriptorInfo *> infoGetters)
{
    if(!m_device || !dstSet) return false;
    uint32_t count = static_cast<uint32_t>(infoGetters.size());
    if(count < 1) return false;

    for (uint32_t i = 0; i < count; ++i) {
        if(!infoGetters[i]) return false;
    }

    std::vector<vk::WriteDescriptorSet> writeinfos;
    std::vector<std::shared_ptr<vk::DescriptorBufferInfo>> bufInfos;
    std::vector<std::shared_ptr<vk::DescriptorImageInfo>> imgInfos;

    for (uint32_t i = 0; i < count; ++i) {
        auto bindinginfo = infoGetters[i]->getDescriptorSetLayoutBinding();
        std::shared_ptr<vk::DescriptorBufferInfo> bufInfo(nullptr);
        std::shared_ptr<vk::DescriptorImageInfo> imgInfo(nullptr);

        if(bindinginfo.descriptorCount < 1){
            SPDLOG_ERROR(__FUNCTION__ "descriptorCount of layout binding is less than 1!");
            bindinginfo.descriptorCount = 1;
        }

        bool hasbufview = false;
        vk::WriteDescriptorSet wrinfo(
                    *dstSet,
                    bindinginfo.binding,
                    0,
                    bindinginfo.descriptorCount,
                    bindinginfo.descriptorType
                    );

        { //  获取DescriptorBufferInfo
            bufInfo = infoGetters[i]->getDescriptorBufferInfo();
            if(bufInfo){
                bufInfos.push_back(bufInfo);
                wrinfo.pBufferInfo = bufInfo.get();
            }
        }

        {  //  获取BufferView
            auto bufview = infoGetters[i]->getBufferView().get();
            if(bufview){
                wrinfo.pTexelBufferView = bufview;
                hasbufview = true;
            }
        }

        { // 获取DescriptorImageInfo
            imgInfo = infoGetters[i]->getDescriptorImageInfo();
            if(imgInfo){
                imgInfos.push_back(imgInfo);
                wrinfo.pImageInfo = imgInfo.get();
            }
        }

        //  基本信息检查
        if(bindinginfo.descriptorType == vk::DescriptorType::eStorageBuffer
                || bindinginfo.descriptorType == vk::DescriptorType::eStorageTexelBuffer
                || bindinginfo.descriptorType == vk::DescriptorType::eUniformBuffer
                || bindinginfo.descriptorType == vk::DescriptorType::eUniformTexelBuffer
                ){
            if(!bufInfo){
                SPDLOG_ERROR(__FUNCTION__
                             " Can not get DescriptorBufferInfo from input"
                             "where binding={}, desctype={}",
                             bindinginfo.binding,
                             static_cast<uint32_t>(bindinginfo.descriptorType)
                             );
                return false;
            }

            if(bindinginfo.descriptorType == vk::DescriptorType::eStorageTexelBuffer
                    || bindinginfo.descriptorType == vk::DescriptorType::eUniformTexelBuffer
                    ){
                if(!hasbufview){
                    SPDLOG_ERROR(__FUNCTION__
                                 " Can not get DescriptorBufferView from input"
                                 " where binding={}, desctype={}",
                                 bindinginfo.binding,
                                 static_cast<uint32_t>(bindinginfo.descriptorType)
                                 );
                    return false;
                }
            }
        }else if(bindinginfo.descriptorType == vk::DescriptorType::eStorageImage
                 || bindinginfo.descriptorType == vk::DescriptorType::eCombinedImageSampler
                 || bindinginfo.descriptorType == vk::DescriptorType::eSampledImage
                 ){
            if(!hasbufview){
                SPDLOG_ERROR(__FUNCTION__
                             " Can not get DescriptorImageInfo from input"
                             " where binding={}, desctype={}",
                             bindinginfo.binding,
                             static_cast<uint32_t>(bindinginfo.descriptorType)
                             );
                return false;
            }
        }else{
            SPDLOG_WARN(__FUNCTION__
                        " Current version did not check"
                        " where binding={},desctype={}",
                        bindinginfo.binding,
                        static_cast<uint32_t>(bindinginfo.descriptorType)
                        );
        }

        writeinfos.push_back(wrinfo);
    }

    m_device->updateDescriptorSets(writeinfos, {});
    bufInfos.clear();
    imgInfos.clear();
    return true;
}

void DescriptorPool::createDescriptorPool(uint32_t maxSetsCount,
                                          uint32_t storageBufSize,
                                          uint32_t uniformTexelBufSize,
                                          uint32_t uniformBufSize,
                                          uint32_t storageTexelBufSize,
                                          uint32_t storageImgSize)
{
    std::vector<vk::DescriptorPoolSize> descriptorPoolSizes;
    if(storageBufSize > 0){
        descriptorPoolSizes.push_back(
                    vk::DescriptorPoolSize(
                        vk::DescriptorType::eStorageBuffer,
                        storageBufSize)
                    );
    }

    if(uniformTexelBufSize > 0){
        descriptorPoolSizes.push_back(
                    vk::DescriptorPoolSize(
                        vk::DescriptorType::eUniformTexelBuffer,
                        uniformTexelBufSize)
                    );
    }

    if(uniformBufSize > 0){
        descriptorPoolSizes.push_back(
                    vk::DescriptorPoolSize(
                        vk::DescriptorType::eUniformBuffer,
                        uniformBufSize)
                    );
    }

    if(storageTexelBufSize > 0){
        descriptorPoolSizes.push_back(
                    vk::DescriptorPoolSize(
                        vk::DescriptorType::eStorageTexelBuffer,
                        storageTexelBufSize)
                    );
    }

    if(storageImgSize > 0){
        descriptorPoolSizes.push_back(
                    vk::DescriptorPoolSize(
                        vk::DescriptorType::eStorageImage,
                        storageImgSize)
                    );
    }

    if(maxSetsCount < 1) maxSetsCount = 1;

    vk::DescriptorPoolCreateInfo descriptorPoolInfo(
                m_descPoolCFlags,
                maxSetsCount, // Max sets
                static_cast<uint32_t>(descriptorPoolSizes.size()),
                descriptorPoolSizes.data());

    auto descPool = std::make_shared<vk::DescriptorPool>();
    auto result = m_device->createDescriptorPool(
                &descriptorPoolInfo, nullptr, descPool.get());

    if(result != vk::Result::eSuccess || !descPool){
        throw std::runtime_error("create vk::DescriptorPool failed!");
    }

    m_descPool = descPool;
}

std::shared_ptr<vk::DescriptorSet>
DescriptorPool::allocateDescriptorSet(std::shared_ptr<vk::DescriptorPool> descpool,
                                      std::shared_ptr<vk::DescriptorSetLayout> layout)
{
    if(!m_device || !descpool || !layout){
        SPDLOG_ERROR(__FUNCTION__ " input with null pointer!");
        return nullptr;
    }

    auto descset = std::make_shared<vk::DescriptorSet>();

    vk::DescriptorSetAllocateInfo descriptorSetAllocateInfo(
                *descpool,
                1, // 创建多少个set
                layout.get());

    m_device->allocateDescriptorSets(&descriptorSetAllocateInfo,
                                     descset.get());

    return descset;
}
